package com.yanti.util;

import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.DeserializationContext;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.deser.std.StdDeserializer;
import com.fasterxml.jackson.databind.module.SimpleModule;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.data.redis.serializer.RedisSerializer;
import org.springframework.data.redis.serializer.SerializationException;

import java.io.IOException;

public class MessageRedisSerializer implements RedisSerializer<Message> {
    private final ObjectMapper objectMapper;

    public MessageRedisSerializer(ObjectMapper objectMapper) {
        this.objectMapper = objectMapper.copy();
        this.objectMapper.registerModule(new SimpleModule()
                .addDeserializer(Message.class, new MessageDeserializer())
        );
    }

    @Override
    public byte[] serialize(Message message) throws SerializationException {
        try {
            return objectMapper.writeValueAsBytes(message);
        } catch (JsonProcessingException e) {
            throw new SerializationException("序列化失败", e);
        }
    }

    @Override
    public Message deserialize(byte[] bytes) throws SerializationException {
        if (bytes == null || bytes.length == 0) return null;
        try {
            return objectMapper.readValue(bytes, Message.class);
        } catch (IOException e) {
            throw new SerializationException("反序列化失败", e);
        }
    }

    static class MessageDeserializer extends StdDeserializer<Message> {
        public MessageDeserializer() {
            super(Message.class);
        }

        @Override
        public Message deserialize(JsonParser p, DeserializationContext ctxt) throws IOException {
            JsonNode node = p.getCodec().readTree(p);
            String type = node.get("messageType").asText();
            String content = node.get("content").asText();

            return switch (type) {
                case "user" -> new UserMessage(content);
                case "assistant" -> new AssistantMessage(content);
                case "system" -> new SystemMessage(content);
                default -> throw new IllegalArgumentException("未知消息类型: " + type);
            };
        }
    }
}